package com.lw.leetcode.str.c;

import java.util.Map;
import java.util.Stack;
import java.util.TreeMap;

/**
 * Created with IntelliJ IDEA.
 * 726. 原子的数量
 *
 * @author liw
 * @version 1.0
 * @date 2022/6/8 21:23
 */
public class CountOfAtoms {

    public static void main(String[] args) {
        CountOfAtoms test = new CountOfAtoms();

        // H2O
//        String str = "H2O";

        // H2MgO2
//        String str = "Mg(OH)2";

        // H2MgO2
//        String str = "Mg20(O25H2)12";

        // DEHMg2OR
//        String str = "Mg(OH)FDERMg";

        // K4N2O14S4
        String str = "K4(ON(SO3)2)2";

        // K8N4O28S8
//        String str = "((((K4((ON(SO3)2)2)))2))";

        // H2MgNO
//        String str = "Mg(H2O)N";

        String s = test.countOfAtoms(str);
        System.out.println(s);
    }

    private char[] arr;
    private int length;
    private StringBuilder sb = new StringBuilder();

    public String countOfAtoms(String formula) {
        Stack<String> values = new Stack<>();
        Stack<Integer> counts = new Stack<>();
        Stack<String> values2 = new Stack<>();
        Stack<Integer> counts2 = new Stack<>();
        this.arr = formula.toCharArray();
        this.length = formula.length();
        int index = 0;
        int f = 0;
        while (true) {
            Node node = getNext(index);
            int flag = node.flag;
            if (flag == 5) {
                break;
            }
            if (flag == 1) {
                int count = node.count;
                if (f == 2) {
                    counts.add(count);
                } else if (f == 4) {
                    while (!"(".equals(values.peek())) {
                        values2.add(values.pop());
                        counts2.add(counts.pop() * count);
                    }
                    values.pop();
                    while (!counts2.isEmpty()) {
                        counts.add(counts2.pop());
                        values.add(values2.pop());
                    }
                }
            } else if (flag == 2) {
                if (f == 4) {
                    while (!"(".equals(values.peek())) {
                        values2.add(values.pop());
                        counts2.add(counts.pop());
                    }
                    values.pop();
                    while (!counts2.isEmpty()) {
                        counts.add(counts2.pop());
                        values.add(values2.pop());
                    }
                }
                values.add(node.str);
            } else if (flag == 3) {
                values.add(node.op);
            } else {
                if (f == 4) {
                    while (!"(".equals(values.peek())) {
                        values2.add(values.pop());
                        counts2.add(counts.pop());
                    }
                    values.pop();
                    while (!counts2.isEmpty()) {
                        counts.add(counts2.pop());
                        values.add(values2.pop());
                    }
                }
            }
            if (f == 2 && flag != 1) {
                counts.add(1);
            }
            f = flag;
            index = node.index;
        }
        if (f == 2) {
            counts.add(1);
        }
        TreeMap<String, Integer> map = new TreeMap<>();
        while (!values.isEmpty()) {
            String pop = values.pop();
            if ("(".equals(pop)) {
                continue;
            }
            map.put(pop, map.getOrDefault(pop, 0) + counts.pop());
        }
        sb.setLength(0);
        for (Map.Entry<String, Integer> entry : map.entrySet()) {
            sb.append(entry.getKey());
            if (entry.getValue() != 1) {
                sb.append(entry.getValue());
            }
        }
        return sb.toString();
    }

    private Node getNext(int index) {
        if (index == length) {
            return Node.getEnd();
        }
        char c = arr[index];
        if (c == '(') {
            return Node.getleft(++index);
        }
        if (c == ')') {
            return Node.getRight(++index);
        }
        if (c <= '9') {
            int count = c - '0';
            if (index == length - 1 || arr[index + 1] > '9' || arr[index + 1] < '0') {
                return Node.getInt(count, index + 1);
            }
            index++;
            while (index < length && arr[index] <= '9' && arr[index] >= '0') {
                count = count * 10 + arr[index] - '0';
                index++;
            }
            return Node.getInt(count, index);
        }
        if (index == length - 1 || arr[index + 1] < 'a') {
            return Node.getStr(String.valueOf(c), index + 1);
        }
        sb.setLength(0);
        sb.append(c);
        index++;
        while (index < length && arr[index] >= 'a') {
            sb.append(arr[index]);
            index++;
        }
        return Node.getStr(sb.toString(), index);
    }

    private static class Node {
        private int count;
        private String str;
        private String op;
        private int flag;
        private int index;

        private Node() {
        }

        private static Node getEnd() {
            Node node = new Node();
            node.flag = 5;
            return node;
        }

        private static Node getInt(int count, int index) {
            Node node = new Node();
            node.flag = 1;
            node.index = index;
            node.count = count;
            return node;
        }

        private static Node getStr(String str, int index) {
            Node node = new Node();
            node.flag = 2;
            node.index = index;
            node.str = str;
            return node;
        }

        private static Node getleft(int index) {
            Node node = new Node();
            node.flag = 3;
            node.index = index;
            node.op = "(";
            return node;
        }

        private static Node getRight(int index) {
            Node node = new Node();
            node.flag = 4;
            node.index = index;
            node.op = ")";
            return node;
        }
    }

}
